#!/usr/bin/env python3
"""
Stable Video Diffusion 模型下载脚本

从 HuggingFace Hub 下载 SVD 模型和相关组件。
"""

import os
import sys
import yaml
import logging
from pathlib import Path
from huggingface_hub import hf_hub_download, snapshot_download

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class ModelDownloader:
    def __init__(self, config_path='config/config.yaml'):
        """
        初始化模型下载器
        :param config_path: 配置文件路径
        """
        self.config = self._load_config(config_path)
        if not self.config:
            logger.error("配置文件加载失败，请先运行 setup_environment.py")
            return
            
        self.model_name = self.config['model']['name']
        self.cache_dir = self.config['model']['cache_dir']
        self.use_safetensors = self.config['model']['use_safetensors']

    def _load_config(self, config_path):
        """加载YAML配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            logger.error(f"配置文件未找到: {config_path}")
            return None
        except Exception as e:
            logger.error(f"加载配置文件时出错: {e}")
            return None

    def download_model(self):
        """
        下载或验证指定的SVD模型.
        snapshot_download is idempotent and will only download missing files.
        """
        if not self.config:
            return

        logger.info("="*50)
        logger.info(f"开始验证/下载模型: {self.model_name}")
        logger.info(f"模型将保存到缓存目录: {self.cache_dir}")
        logger.info("="*50)
        
        # 确保缓存目录存在
        Path(self.cache_dir).mkdir(exist_ok=True)

        try:
            # 使用 snapshot_download 下载整个仓库。
            # 这会自动处理缓存，如果文件已存在则不会重复下载。
            # 我们不再需要 local_dir，让 huggingface_hub 管理缓存结构。
            snapshot_download(
                repo_id=self.model_name,
                cache_dir=self.cache_dir,
                local_dir_use_symlinks=False,
                allow_patterns=["*.json", "*.safetensors" if self.use_safetensors else "*.bin"],
                ignore_patterns=["*.ckpt", "*.pt", "*.png"],
                token=os.environ.get("HUGGINGFACE_TOKEN") # 从环境变量读取token
            )
            logger.info(f"\n✅ 模型 {self.model_name} 验证/下载成功！")
            logger.info(f"模型文件位于缓存目录: {self.cache_dir}")
            return True

        except Exception as e:
            logger.error(f"\n❌ 模型下载失败: {e}")
            logger.error("请检查：")
            logger.error("1. 网络连接是否正常")
            logger.error("2. HuggingFace Hub 是否可访问")
            logger.error("3. 如果是私有模型，是否设置了 HUGGINGFACE_TOKEN 环境变量")
            return False

if __name__ == "__main__":
    downloader = ModelDownloader()
    if downloader.config:
        if not downloader.download_model():
            sys.exit(1)
